﻿/*
 * PsParser.cpp
 *
 *  Created on: 2017年4月18日
 *      Author: chuanjiang.zh
 */

#include "PsParser.h"
#include "TStringUtil.h"
#include "CLog.h"
#include "mpeg.h"
#include "TFileUtil.h"


namespace av
{

#ifndef WIN32
typedef unsigned char BYTE;
typedef unsigned short WORD;

#define MAKEWORD(a,b) ((WORD) (((BYTE) (a)) | ((WORD) ((BYTE) (b))) << 8))

#endif //WIN32


static uint8_t START_CODE[] = {0, 0, 1};
static unsigned int IMKH = MKTAG('I','M','K', 'H');

static size_t	SYNC_BYTES_LENGTH = 4;	/// including start code

static size_t	MIN_HEADER_SIZE = 6;

static size_t	MAX_PS_LENGTH = 0x100000;
static size_t	MAX_PES_LENGTH = 0xFFFF;
static size_t	MAX_ES_LENGTH = 0x100000;


enum PSStatus
{
	ps_padding, //未知状态
	ps_ps,      //ps状态
	ps_sh,
	ps_psm,
	ps_pes,
	ps_pes_video,
	ps_pes_audio,
	ps_pes_last,
	ps_pes_lost, //pes包丢失数据（udp传输的rtp丢包）
	ps_pes_jump   //需要跳过的pes packet
};

#pragma pack(1)

//ps header
typedef struct ps_header {
	unsigned char pack_start_code[4];  //'0x000001BA'

	unsigned char system_clock_reference_base21 : 2;
	unsigned char marker_bit : 1;
	unsigned char system_clock_reference_base1 : 3;
	unsigned char fix_bit : 2;    //'01'

	unsigned char system_clock_reference_base22;

	unsigned char system_clock_reference_base31 : 2;
	unsigned char marker_bit1 : 1;
	unsigned char system_clock_reference_base23 : 5;

	unsigned char system_clock_reference_base32;
	unsigned char system_clock_reference_extension1 : 2;
	unsigned char marker_bit2 : 1;
	unsigned char system_clock_reference_base33 : 5; //system_clock_reference_base 33bit

	unsigned char marker_bit3 : 1;
	unsigned char system_clock_reference_extension2 : 7; //system_clock_reference_extension 9bit

	unsigned char program_mux_rate1;

	unsigned char program_mux_rate2;
	unsigned char marker_bit5 : 1;
	unsigned char marker_bit4 : 1;
	unsigned char program_mux_rate3 : 6;

	unsigned char pack_stuffing_length : 3;
	unsigned char reserved : 5;
}ps_header_t;  //14

//system header
typedef struct sh_header
{
	unsigned char system_header_start_code[4]; //32

	unsigned char header_length[2];            //16 uimsbf

	uint32_t marker_bit1 : 1;   //1  bslbf
	uint32_t rate_bound : 22;   //22 uimsbf
	uint32_t marker_bit2 : 1;   //1 bslbf
	uint32_t audio_bound : 6;   //6 uimsbf
	uint32_t fixed_flag : 1;    //1 bslbf
	uint32_t CSPS_flag : 1;     //1 bslbf

	uint16_t system_audio_lock_flag : 1;  // bslbf
	uint16_t system_video_lock_flag : 1;  // bslbf
	uint16_t marker_bit3 : 1;             // bslbf
	uint16_t video_bound : 5;             // uimsbf
	uint16_t packet_rate_restriction_flag : 1; //bslbf
	uint16_t reserved_bits : 7;                //bslbf
	unsigned char reserved[6];
}sh_header_t; //18

//program stream map header
typedef struct psm_header {
	unsigned char promgram_stream_map_start_code[4];

	unsigned char program_stream_map_length[2];

	unsigned char program_stream_map_version : 5;
	unsigned char reserved1 : 2;
	unsigned char current_next_indicator : 1;

	unsigned char marker_bit : 1;
	unsigned char reserved2 : 7;

	unsigned char program_stream_info_length[2];
	unsigned char elementary_stream_map_length[2];

	
}psm_header_t; //12

typedef struct psm_stream_info {
	unsigned char stream_type;
	unsigned char elementary_stream_id;
	unsigned char elementary_stream_info_length[2];
	unsigned char crc_32[4];
	// stream info data
} psm_stream_info_t;

			   
typedef struct pes_header
{
	unsigned char pes_start_code_prefix[3];
	unsigned char stream_id;
	unsigned char pes_packet_length[2];
}pes_header_t; //6
			   
typedef struct optional_pes_header {
	unsigned char original_or_copy : 1;
	unsigned char copyright : 1;
	unsigned char data_alignment_indicator : 1;
	unsigned char PES_priority : 1;
	unsigned char PES_scrambling_control : 2;
	unsigned char fix_bit : 2;

	unsigned char PES_extension_flag : 1;
	unsigned char PES_CRC_flag : 1;
	unsigned char additional_copy_info_flag : 1;
	unsigned char DSM_trick_mode_flag : 1;
	unsigned char ES_rate_flag : 1;
	unsigned char ESCR_flag : 1;
	unsigned char PTS_DTS_flags : 2;

	unsigned char PES_header_data_length;
}optional_pes_header_t;

#pragma pack()




static MediaCodec getCodecFromStreamType(int stream_type)
{
	if (stream_type == 0x90)
	{
		return MEDIA_CODEC_G711A;
	}
	else if (stream_type == 0x91)
	{
		return MEDIA_CODEC_G711U;
	}
	else if (stream_type == 0x92)
	{
		return MEDIA_CODEC_G722;
	}
	else if (stream_type == STREAM_TYPE_VIDEO_MPEG2)
	{
		return MEDIA_CODEC_MPEG2;
	}
	else if (stream_type == STREAM_TYPE_AUDIO_MPEG2)
	{
		return MEDIA_CODEC_MP3;
	}
	else if (stream_type == STREAM_TYPE_VIDEO_MPEG4)
	{
		return MEDIA_CODEC_MPEG4;
	}
	else if (stream_type == STREAM_TYPE_AUDIO_AAC)
	{
		return MEDIA_CODEC_AAC;
	}
	else if (stream_type == STREAM_TYPE_VIDEO_H264)
	{
		return MEDIA_CODEC_H264;
	}
	else if (stream_type == STREAM_TYPE_AUDIO_AC3)
	{
		return MEDIA_CODEC_AC3;
	}
	else if (stream_type == STREAM_TYPE_VIDEO_HEVC)
	{
		return MEDIA_CODEC_HEVC;
	}

	return MEDIA_CODEC_NONE;
}


PsParser::PsParser():
		m_imkh(),
		m_lostCount()
{
	m_buffer.ensure(1024 * 1024);

	m_pktBuffer.ensure(1024 * 800);
}

PsParser::~PsParser()
{
}


bool PsParser::parse(const uint8_t* buf, size_t length, Packet& pkt)
{
	if (buf && (length > 0))
	{
		m_buffer.write(buf, length);
	}

	bool found = false;
	while (!found)
	{
		uint8_t* data = m_buffer.getReadPtr();
		size_t size = m_buffer.readable();

		if (size < MIN_HEADER_SIZE)
		{
			break;
		}

		size_t startIndex = findHeaderStartCode(data, 0, size);
		if (startIndex == (size_t)-1)
		{
			if (size > MAX_PES_LENGTH)
			{
				CLog::warning("PsParser. found data without start code.\n");
				m_buffer.clear();
			}

			break;
		}
		else if (startIndex != 0)
		{
			CLog::warning("PsParser. skip data with size %d\n", startIndex);
			m_buffer.skip(startIndex);

			data = m_buffer.getReadPtr();
			size = m_buffer.readable();
		}

		size_t len = parsePacket(data, size, found, pkt);
		if (len == -1)
		{
			// error
			m_buffer.skip(MIN_HEADER_SIZE);
		}
		else if (len == 0)
		{
			// not enough data
			break;
		}
		else
		{
			// consumed
			m_buffer.skip(len);
		}
	}

	if (found)
	{
		pkt = m_outPkt;
	}
	return found;
}

void PsParser::clear()
{
	m_buffer.clear();

	m_pktBuffer.clear();
	m_pkt = Packet();
	m_outPkt = Packet();

	m_imkh = false;
}

const MediaFormat& PsParser::getFormat()
{
	return m_format;
}

size_t PsParser::parsePacket(uint8_t* data, size_t length, bool& found, Packet& packet)
{
	if (!startWithCode(data, length))
	{
		return -1;
	}

	uint8_t code = data[3];
	if (code == 0xBA)
	{
		return parsePackHeader(data, length);
	}
	else if (code == 0xBB)
	{
		return parseSystemHeader(data, length);
	}
	else if (code == 0xBC)
	{
		return parsePsmHeader(data, length);
	}
	else if (code >= 0xC0 && code <= 0xDF)
	{
		return parsePesPacket(MEDIA_TYPE_AUDIO, data, length, found, packet);
	}
	else if (code >= 0xE0 && code <= 0xEF)
	{
		return parsePesPacket(MEDIA_TYPE_VIDEO, data, length, found, packet);
	}
	else if (code >= 0xBD && code <= 0xBF)
	{
		return parseJump(data, length);
	}
	else if (code >= 0xF0 && code <= 0xF8)
	{
		return parseJump(data, length);
	}

	return -1;
}

size_t PsParser::parsePackHeader(uint8_t* data, size_t length)
{
	if (length < sizeof(ps_header_t))
	{
		return 0;
	}

	ps_header_t* header = (ps_header_t*)data;
	size_t wanted = sizeof(ps_header_t) + header->pack_stuffing_length;

	if (length < wanted)
	{
		return 0;
	}

	return wanted;
}

size_t PsParser::parseSystemHeader(uint8_t* data, size_t length)
{
	if (length < sizeof(sh_header_t))
	{
		return 0;
	}

	sh_header_t* header = (sh_header_t*)data;
	unsigned short header_length = MAKEWORD(header->header_length[1], header->header_length[0]);
	size_t wanted = SYNC_BYTES_LENGTH + 2 + header_length;

	if (length < wanted)
	{
		return 0;
	}

	return wanted;
}

size_t PsParser::parsePsmHeader(uint8_t* data, size_t length)
{
	if (length < sizeof(psm_header_t))
	{
		return 0;
	}

	psm_header_t* header = (psm_header_t*)data;
	uint16_t stream_map_length = MAKEWORD(header->program_stream_map_length[1], header->program_stream_map_length[0]);
	uint16_t stream_info_length = MAKEWORD(header->program_stream_info_length[1], header->program_stream_info_length[0]);
	uint16_t es_map_length = MAKEWORD(header->elementary_stream_map_length[1], header->elementary_stream_map_length[0]);
	size_t wanted = SYNC_BYTES_LENGTH + 2 + stream_map_length;

	if (length < wanted)
	{
		return 0;
	}

	Packet pkt;
	pkt.data = data;
	pkt.length = wanted;

	parsePsmHeader(pkt);

	//uint8_t* esdata = data + sizeof(psm_header_t);
	//if (es_map_length > 0)
	//{
	//	psm_stream_info_t* info = (psm_stream_info_t*)esdata;
	//	uint16_t info_length = MAKEWORD(info->elementary_stream_info_length[1], info->elementary_stream_info_length[0]);
	//	uint8_t* info_data = esdata + sizeof(psm_stream_info_t);

	//	parseStreamInfo(info->elementary_stream_id, info->stream_type, info_data, info_length);

	//	if (es_map_length > (sizeof(psm_stream_info_t) + info_length))
	//	{
	//		esdata += sizeof(psm_stream_info_t) + info_length;
	//		es_map_length -= sizeof(psm_stream_info_t) + info_length;

	//		uint16_t info_length = MAKEWORD(info->elementary_stream_info_length[1], info->elementary_stream_info_length[0]);
	//		uint8_t* info_data = esdata + sizeof(psm_stream_info_t);

	//		parseStreamInfo(info->elementary_stream_id, info->stream_type, info_data, info_length);
	//	}
	//}

	return wanted;
}

size_t PsParser::parseJump(uint8_t* data, size_t length)
{
	size_t startIndex = findHeaderStartCode(data, SYNC_BYTES_LENGTH, length);
	if (startIndex == (size_t)-1)
	{
		return 0;
	}
	return startIndex;
}

size_t PsParser::parsePesPacket(int streamType, uint8_t* data, size_t length, bool& found, Packet& packet)
{
	if (length < sizeof(pes_header_t))
	{
		return 0;
	}

	pes_header_t* header = (pes_header_t*)data;
	unsigned short pes_packet_length = MAKEWORD(header->pes_packet_length[1], header->pes_packet_length[0]);
	size_t wanted = SYNC_BYTES_LENGTH + 2 + pes_packet_length;

	if (length < wanted)
	{
		return 0;
	}

	if (pes_packet_length == 0)
	{
		CLog::info("PsParser. pes_packet_length == 0\n");

		size_t startIndex = findHeaderStartCode(data, 6, length);
		if (startIndex == (size_t)-1)
		{
			return 0;
		}

		uint8_t* payload = data + sizeof(pes_header_t);
		int size = startIndex - sizeof(pes_header_t);

		m_pktBuffer.getCurBuffer().write(payload, size);

		return startIndex;
	}

	optional_pes_header_t* pes_header = (optional_pes_header_t*)(data + 6);

	int64_t pts = 0;
	bool hasPts = false;

	if ((pes_header->PES_header_data_length >= 5) && (pes_header->PTS_DTS_flags != 0))
	{
		/// parse ts
		int index = 9;
		int64_t b30 = (data[index] & 0x0E) >> 1;
		index++;
		int64_t b22 = data[index];
		index++;
		int64_t b15 = data[index] >> 1;
		index++;
		int64_t b8 = data[index];
		index++;
		int64_t b0 = data[index] >> 1;
		pts = b0 + (b8 << 7) + (b15 << 15) + (b22 << 22) + (b30 << 30);
		
		hasPts = true;
		//CLog::debug("pts: %lld\n", pts);
	}

	int hederSize = sizeof(optional_pes_header_t) + pes_header->PES_header_data_length;
	uint8_t* payload = data + sizeof(pes_header_t) + hederSize;
	int size = pes_packet_length - hederSize;


	/// check lost
	if (length < (wanted + SYNC_BYTES_LENGTH))
	{
		return 0;
	}

	uint8_t* nextData = data + wanted;
	if (!startWithCode(nextData, SYNC_BYTES_LENGTH))
	{
		CLog::warning("PsParser. detect packet lost. count:%d\n", m_lostCount);
		logData("lost packet.", payload, size);

		m_lostCount ++;

		size_t offset = sizeof(pes_header_t) + sizeof(optional_pes_header_t);
		size_t startIndex = findHeaderStartCode(data, offset, length);
		if (startIndex == (size_t)-1)
		{
			/// current is complete
		}
		else
		{
			wanted = startIndex;
			size = (data + startIndex) - payload;
		}
	}

	//
	if (m_pkt.type != streamType)
	{
		found = flushPacket();
	}
	else if (hasPts)
	{
		if (m_pkt.ts != pts)
		{
			found = flushPacket();
		}
	}

	m_pktBuffer.getCurBuffer().write(payload, size);
	m_pkt.type = streamType;

	if (hasPts)
	{
		m_pkt.ts = pts;
	}

	//comn::FileUtil::write(payload, size, "pes.h264", true);

	//logData("pes", payload, size);

	return wanted;
}

bool PsParser::parsePacket(Packet& packet)
{
	bool found = false;
	if (!startWithCode(packet.data, packet.length))
	{
		ByteBuffer& buffer = m_pktBuffer.getCurBuffer();
		if (buffer.size() == 0)
		{
			if (packet.length > sizeof(IMKH))
			{
				unsigned int startcode = MKTAG(packet.data[0], packet.data[1], packet.data[2], packet.data[3]);
				m_imkh = (startcode == IMKH);
			}
		}
		return flushPacket();
	}

	uint8_t code = packet.data[3];

	if (code == 0xBA)
	{
		found = flushPacket();
		parsePackHeader(packet);
	}
	else if (code == 0xBB)
	{
		found = flushPacket();
		parseSystemHeader(packet);
	}
	else if (code == 0xBC)
	{
		found = flushPacket();
		parsePsmHeader(packet);
	}
	else if (code <= 0xBF)
	{
		found = flushPacket();
	}
	else if (code >= 0xC0 && code <= 0xDF)
	{
		// audio stream
		packet.type = MEDIA_TYPE_AUDIO;
		found = parsePESPacket(packet);
	}
	else if (code >= 0xE0 && code <= 0xED)
	{
		// video stream
		packet.type = MEDIA_TYPE_VIDEO;
		found = parsePESPacket(packet);
	}
	else
	{
		// ignore
	}

	return found;
}

bool PsParser::parsePESPacket(Packet& packet)
{

	uint16_t pktLength = MAKEWORD(packet.data[5], packet.data[4]);
	size_t hdrLength = packet.data[5 + 3];
	size_t offset = sizeof(START_CODE) + 3 + 3 + hdrLength;

	bool hasPts = false;
	int64_t pts = 0;

	if (hdrLength >= 5)
	{
		hasPts = true;

		int index = 9;
		int64_t b30 = (packet.data[index] & 0x0E) >> 1;
		index++;
		int64_t b22 = packet.data[index];
		index++;
		int64_t b15 = packet.data[index] >> 1;
		index++;
		int64_t b8 = packet.data[index];
		index++;
		int64_t b0 = packet.data[index] >> 1;
		pts = b0 + (b8 << 7) + (b15 << 15) + (b22 << 22) + (b30 << 30);

		//std::string pktText = comn::StringUtil::format("%d. pktLength:%d,\t hdrLength:%d, \t %02x,%02x,%02x,%02x,%02x,%02x,%02x  pts:%I64d",
		//        s_count, pktLength, hdrLength,
		//            packet.data[9], packet.data[9+1],
		//            packet.data[9+2], packet.data[9+3],
		//            packet.data[9+4], packet.data[9+5],
		//            packet.data[9+6], pts);
		//CLog::debug("ps pkt(%d): %s\n", packet.type, pktText.c_str());
	}

	if (packet.length < offset)
	{
		return false;
	}

	bool found = false;

	uint8_t* data = packet.data + offset;
	uint16_t length = (uint16_t)(packet.length - offset);

	if (hasPts)
	{
		if ((packet.type != m_pkt.type) || (pts != m_pkt.ts))
		{
			found = flushPacket();
		}

		m_pktBuffer.getCurBuffer().write(data, length);
		m_pkt.type = packet.type;
		m_pkt.ts = pts;
	}
	else
	{
		m_pktBuffer.getCurBuffer().write(data, length);
	}

	return found;
}

bool PsParser::parsePackHeader(Packet& packet)
{
	return true;
}

bool PsParser::parseSystemHeader(Packet& packet)
{
	return true;
}

bool PsParser::parsePsmHeader(Packet& packet)
{
	uint8_t* data = packet.data + 8;
	uint16_t stream_info_length = MAKEWORD(data[1], data[0]);
	data += 2 + stream_info_length;
	
	uint16_t stream_map_length = MAKEWORD(data[1], data[0]);
	data += 2;

	{
		uint8_t stream_type = data[0];
		uint8_t stream_id = data[1];
		uint16_t info_length = MAKEWORD(data[3], data[2]);
		data += 4;

		parseStreamInfo(stream_id, stream_type, data, info_length);
		data += info_length;
	}

	if (data < (packet.data + packet.length))
	{
		uint8_t stream_type = data[0];
		uint8_t stream_id = data[1];
		uint16_t info_length = MAKEWORD(data[3], data[2]);
		data += 4;

		parseStreamInfo(stream_id, stream_type, data, info_length);
		data += info_length;
	}

	return true;
}

void PsParser::parseStreamInfo(uint8_t stream_id, uint8_t stream_type, uint8_t* data, int length)
{
	if (stream_id >= 0xC0 && stream_id <= 0xDF)
	{
		m_format.m_audioCodec = getCodecFromStreamType(stream_type);

		m_format.m_sampleRate = 8000;
		m_format.m_channels = 1;
		if (m_format.m_audioCodec == av::MEDIA_CODEC_G722)
		{
			m_format.m_sampleRate = 16000;
		}
	}
	else if (stream_id >= 0xE0 && stream_id <= 0xED)
	{
		m_format.m_codec = getCodecFromStreamType(stream_type);
	}
}

bool PsParser::startWithCode(const uint8_t* buffer, size_t length)
{
	return (length > sizeof (START_CODE)) &&
			(buffer[0] == 0) &&
			(buffer[1] == 0) &&
			(buffer[2] == 1);
}

size_t PsParser::findStartCode(const uint8_t* buffer, size_t length)
{
    if (length < sizeof(START_CODE))
    {
        return -1;
    }

    for (size_t i = 0; i < (length - sizeof(START_CODE)); ++ i)
    {
        if ((buffer[i] == 0) && (buffer[i+1] == 0) && (buffer[i+2] == 1))
        {
            return i;
        }
    }
    return -1;
}

size_t PsParser::findStartCode(const uint8_t* buffer, size_t offset, size_t length)
{
    if (length <= offset)
    {
        return -1;
    }

    size_t idx = findStartCode(buffer + offset, length - offset);
    if (idx != (size_t)-1)
    {
        idx += offset;
    }
    return idx;
}

size_t PsParser::findHeaderStartCode(const uint8_t* buffer, size_t offset, size_t length)
{
    size_t idx = findStartCode(buffer, offset, length);
    while (idx != (size_t)-1)
    {
        uint8_t type = buffer[idx + sizeof(START_CODE)];
		if ((type >= 0xBA) && (type <= 0xF8))
        {
            break;
        }
        else
        {
            idx = findStartCode(buffer, idx + sizeof(START_CODE), length);
        }
    }
    return idx;
}

bool PsParser::flushPacket()
{
	ByteBuffer& buffer = m_pktBuffer.getCurBuffer();
	if (buffer.size() == 0)
	{
		return false;
	}

	m_outPkt = m_pkt;
	m_outPkt.data = buffer.data();
	m_outPkt.length = buffer.size();

	m_pktBuffer.switchBuffer();
	m_pktBuffer.getCurBuffer().clear();
	m_pkt = Packet();

	return true;
}

void PsParser::logData(const char* name, const uint8_t* data, int size)
{
	int minSize = std::min(size, 16);
	std::string hex = comn::StringCast::toHexGroupString(data, minSize);
	CLog::debug("%s. size:%d, %s\n", name, size, hex.c_str());
}



} /* namespace av */
